import backtrader as bt
import pandas as pd
import math


class BollBand(bt.Indicator):
    lines = ('upper', 'middle', 'lower',)
    params = (('maperiod', 20),
              ('period', 3),
              ('multiple', 2),)

    # 与价格在同一张图
    plotinfo = dict(subplot=False)

    def __init__(self):
        ema = bt.ind.EMA(self.data, period=self.p.maperiod)
        # 计算上中下轨线
        self.lines.middle = bt.ind.EMA(ema, period=self.p.period)
        self.pow = bt.ind.SMA(pow(self.data - self.lines.middle, 2), period=self.p.period)
        super(BollBand, self).__init__()

    def next(self):
        vart = math.sqrt(self.pow[0])
        self.lines.upper[0] = self.lines.middle[0] + self.p.multiple * vart
        self.lines.lower[0] = self.lines.middle[0] - self.p.multiple * vart


class BollBB(bt.Indicator):
    lines = ('bb', 'bb_ma',)
    params = (('period', 6),
              ('boll_period', 20),
              ('boll_multiple', 2),)

    plotinfo = dict(subplot=True)

    def __init__(self):
        boll = BollBand(self.data, maperiod=self.p.boll_period, multiple=self.p.boll_multiple)
        self.lines.bb = (self.data - boll.lower) / (boll.upper - boll.lower) * 100
        self.lines.bb_ma = bt.ind.SMA(self.lines.bb, period=self.p.period)
        super(BollBB, self).__init__()


class BollWidth(bt.Indicator):
    lines = ('width',)
    params = (('boll_period', 20),
              ('boll_multiple', 2),)

    plotinfo = dict(subplot=True)

    def __init__(self):
        boll = BollBand(self.data, maperiod=self.p.boll_period, multiple=self.p.boll_multiple)
        self.lines.width = (boll.upper - boll.lower) / boll.middle * 100
        super(BollWidth, self).__init__()


class AtrLoss(bt.Indicator):
    lines = ('atr_loss',)
    params = (('period', 21),
              ('multiple', 1.5))

    plotinfo = dict(subplot=False)

    def __init__(self):
        atr = bt.ind.ATR(self.data, period=self.p.period)
        self.atr_loss = self.data.low - atr * self.p.multiple
        self.lines.atr_loss = bt.ind.MaxN(self.atr_loss, period=self.p.period)
        super(AtrLoss, self).__init__()

